from .backbone import MAE_ViT_Cifar
from .mocov2 import MoCoV2_Model
from .simsiam import SimSiam
from .densecl import DenseCL
from .simclr import SimCLR
from .pcl import PCL

def get_model(args):
    if args.model == 'mocov2':
        return MoCoV2_Model(backbone=args.backbone,out_dim=args.out_dim, queue_size=args.queue_size, 
        momentum=args.queue_momentum, temperature=args.temperature)
    elif args.model == 'simsiam':
        return SimSiam(backbone=args.backbone, feat_dim=args.feat_dim,num_proj_layers=args.num_proj_layers)
    elif args.model == 'mae_vit_cifar':
        return MAE_ViT_Cifar(mask_ratio=args.mask_ratio)
    elif args.model == "densecl":
        return DenseCL(backbone=args.backbone, feature_dim=args.feature_dim, K=args.K, m=args.m, T=args.T, 
        loss_lambda=args.loss_lambda, num_grid=args.num_grid)
    elif args.model == "simclr":
        return SimCLR(backbone=args.backbone,out_dim=args.out_dim)
    elif args.model == "pcl":
        return PCL(backbone=args.backbone)
    else:
        raise NotImplementedError
